# init_db.py - 数据库初始化脚本
import sqlite3

# 初始化数据库
def init_database():
    try:
        # 连接数据库（如果不存在则创建）
        conn = sqlite3.connect('electric_meter.db')
        cursor = conn.cursor()
        
        # 创建用电数据表
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS electricity_data (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT,
            voltage REAL,
            current REAL,
            power REAL,
            energy REAL
        )
        ''')
        
        # 创建用电模式分析结果表（用于存储历史分析数据）
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS analysis_results (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT,
            status TEXT,
            suggestion TEXT,
            confidence REAL
        )
        ''')
        
        # 创建系统配置表
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS system_config (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            key TEXT UNIQUE,
            value TEXT
        )
        ''')
        
        # 插入默认配置
        default_configs = [
            ('threshold_voltage_high', '240.0'),
            ('threshold_voltage_low', '200.0'),
            ('threshold_current_high', '10.0'),
            ('threshold_power_high', '2000.0'),
            ('sample_interval', '5.0'),
            ('anomaly_detection_enabled', 'true')
        ]
        
        # 插入或更新配置
        for key, value in default_configs:
            cursor.execute(
                "INSERT OR REPLACE INTO system_config (key, value) VALUES (?, ?)",
                (key, value)
            )
        
        # 提交更改
        conn.commit()
        
        print("数据库初始化成功！")
        print("- 创建了 electricity_data 表用于存储用电数据")
        print("- 创建了 analysis_results 表用于存储分析结果")
        print("- 创建了 system_config 表用于存储系统配置")
        print("- 插入了默认配置值")
        
    except sqlite3.Error as e:
        print(f"数据库初始化失败: {e}")
    finally:
        # 关闭数据库连接
        if conn:
            conn.close()

# 检查数据库是否已初始化
def check_database():
    try:
        conn = sqlite3.connect('electric_meter.db')
        cursor = conn.cursor()
        
        # 检查表格是否存在
        cursor.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name='electricity_data'"
        )
        if cursor.fetchone():
            print("数据库已存在并初始化。")
            
            # 查询一些统计信息
            cursor.execute("SELECT COUNT(*) FROM electricity_data")
            count = cursor.fetchone()[0]
            print(f"- electricity_data 表中有 {count} 条记录")
            
            cursor.execute("SELECT MIN(timestamp), MAX(timestamp) FROM electricity_data")
            time_range = cursor.fetchone()
            if time_range[0]:
                print(f"- 数据时间范围: {time_range[0]} 至 {time_range[1]}")
        else:
            print("数据库未初始化，将进行初始化。")
            init_database()
            
    except sqlite3.Error as e:
        print(f"检查数据库时出错: {e}")
    finally:
        if conn:
            conn.close()

# 清空数据表（用于测试或重置）
def clear_database():
    try:
        conn = sqlite3.connect('electric_meter.db')
        cursor = conn.cursor()
        
        # 清空数据表但保留表结构
        cursor.execute("DELETE FROM electricity_data")
        cursor.execute("DELETE FROM analysis_results")
        
        # 重置自增计数器
        cursor.execute("DELETE FROM sqlite_sequence WHERE name='electricity_data'")
        cursor.execute("DELETE FROM sqlite_sequence WHERE name='analysis_results'")
        
        conn.commit()
        print("数据库已清空！")
        
    except sqlite3.Error as e:
        print(f"清空数据库时出错: {e}")
    finally:
        if conn:
            conn.close()

# 生成示例数据（用于演示）
def generate_sample_data(count=100):
    import random
    import time
    from datetime import datetime, timedelta
    
    try:
        conn = sqlite3.connect('electric_meter.db')
        cursor = conn.cursor()
        
        # 清空现有数据
        cursor.execute("DELETE FROM electricity_data")
        cursor.execute("DELETE FROM sqlite_sequence WHERE name='electricity_data'")
        
        # 生成示例数据
        now = datetime.now()
        for i in range(count):
            # 生成时间戳（从现在往前推）
            timestamp = (now - timedelta(minutes=i*5)).strftime('%Y-%m-%d %H:%M:%S')
            
            # 获取小时数用于模拟不同时间段的用电模式
            hour = (now - timedelta(minutes=i*5)).hour
            
            # 模拟不同时间段的用电模式
            if 0 <= hour < 6:
                # 夜间低用电模式
                voltage = 220.0 + random.uniform(-5, 5)
                current = 0.3 + random.uniform(0, 0.3)
            elif 6 <= hour < 12:
                # 早晨中等用电模式
                voltage = 220.0 + random.uniform(-5, 5)
                current = 0.8 + random.uniform(0, 0.8)
            elif 12 <= hour < 18:
                # 中午中等用电模式
                voltage = 220.0 + random.uniform(-5, 5)
                current = 0.6 + random.uniform(0, 0.6)
            else:
                # 晚上高用电模式
                voltage = 220.0 + random.uniform(-5, 5)
                current = 1.2 + random.uniform(0, 1.2)
            
            # 计算功率和用电量
            power = voltage * current
            energy = power / 12.0  # 转换为kWh（假设每5分钟一个数据点）
            
            # 插入数据
            cursor.execute(
                "INSERT INTO electricity_data (timestamp, voltage, current, power, energy) VALUES (?, ?, ?, ?, ?)",
                (timestamp, voltage, current, power, energy)
            )
        
        conn.commit()
        print(f"已生成 {count} 条示例数据！")
        
    except sqlite3.Error as e:
        print(f"生成示例数据时出错: {e}")
    finally:
        if conn:
            conn.close()

if __name__ == '__main__':
    # 命令行参数解析
    import sys
    
    if len(sys.argv) > 1:
        if sys.argv[1] == 'check':
            check_database()
        elif sys.argv[1] == 'init':
            init_database()
        elif sys.argv[1] == 'clear':
            clear_database()
        elif sys.argv[1] == 'sample':
            count = 100
            if len(sys.argv) > 2:
                try:
                    count = int(sys.argv[2])
                except ValueError:
                    pass
            generate_sample_data(count)
        else:
            print("用法: python init_db.py [check|init|clear|sample]")
            print("  check    - 检查数据库状态")
            print("  init     - 初始化数据库")
            print("  clear    - 清空数据库")
            print("  sample   - 生成示例数据")
    else:
        # 默认行为：检查数据库状态，如果未初始化则进行初始化
        check_database()